import json
import urllib.request
import urllib.parse

def queue_prompt(prompt, client_id, server_address):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request("http://{}/prompt".format(server_address), data=data)
    return json.loads(urllib.request.urlopen(req).read())


def get_image(filename, subfolder, folder_type, server_address):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response:
        return response.read()


def get_history(prompt_id, server_address):
    with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
        return json.loads(response.read())


def get_images(ws, prompt, client_id, server_address, savePath):
    prompt_id = queue_prompt(prompt, client_id, server_address)['prompt_id']
    output_images = {}
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            # print(message)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    history = get_history(prompt_id, server_address)[prompt_id]
                    if not history['outputs']:
                        for messages in history["status"]:
                            message = history["status"]["messages"][2]
                            return [history["status"][messages], message[1]["exception_message"],
                                    message[1]["exception_type"]]
                    else:
                        for node_id in history['outputs']:
                            node_output = history['outputs'][node_id]
                            imageName = ""
                            if 'images' in node_output:
                                for image in node_output['images']:
                                    imageName = image['filename']

                    yield {"node": None, "imageName": savePath + imageName}
                    break  # Execution is done
                else:
                    yield data

        else:
            # If you want to be able to decode the binary stream for latent previews, here is how you can do it:
            # bytesIO = BytesIO(out[8:])
            # preview_image = Image.open(bytesIO) # This is your preview in PIL image format, store it in a global
            continue  # previews are binary data


